package com.sunyw.xyz.api;

import com.alibaba.fastjson.JSON;
import com.beat.chatgpt.model.UnifyChatCompletion;
import com.beat.chatgpt.model.UnifyChoice;
import com.beat.chatgpt.model.UnifyCompletion;
import com.beat.chatgpt.model.UnifyStreamMessage;
import com.beat.chatgpt.provider.openai.dto.completions.CompletionResponse;
import com.sunyw.xyz.model.ChatCompletionDTO;
import com.sunyw.xyz.model.ChatModel;
import com.sunyw.xyz.model.Result;
import com.sunyw.xyz.service.OpenAIService;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.io.IOUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.InputStreamResource;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import org.springframework.web.servlet.mvc.method.annotation.StreamingResponseBody;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.*;
import java.nio.charset.StandardCharsets;
import java.util.Date;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;


/**
 * 名称: XX定义
 * 功能: <功能详细描述>
 * 方法: <方法简述-方法描述>
 * 版本: 1.0
 * 作者: sunyw
 * 说明: 说明描述
 * 时间: 2023/02/07 12:49
 */
@Slf4j
@RestController
@RequestMapping("/api")
public class ChatGptApi {

    @Autowired
    private AsyncService asyncService;

    @Autowired
    private OpenAIService openAIService;

    private final ExecutorService executorService = Executors.newCachedThreadPool();

//    @PostMapping("/question")
//    public Server test(@RequestBody Server server) {
//
//        return server;
//    }


    @PostMapping("/question")
    public void question(HttpServletRequest request) {
        long l = System.currentTimeMillis();
        asyncService.listen(getBody(request));
        log.info("消息接收结束:耗时:[{}]毫秒", (System.currentTimeMillis() - l));
    }

    @PostMapping("/chat")
    public Result<String> chat(@RequestBody ChatModel chatModel) {
        return openAIService.completionsHandleEx(chatModel.getPrompt());
    }


    @PostMapping("/chat/completion/stream")
    public ResponseEntity<InputStreamResource> chatStream(@RequestBody ChatCompletionDTO chatCompletion) {

        HttpHeaders headers = new HttpHeaders();
        headers.setContentType(MediaType.TEXT_EVENT_STREAM);
        headers.setCacheControl("no-cache");

        InputStream inputStream=openAIService.chatCompletionStream(chatCompletion);
        InputStreamResource resource = new InputStreamResource(inputStream);

        return new ResponseEntity<InputStreamResource>(resource, headers, HttpStatus.OK);
    }

    @PostMapping("/chat/completion")
    public Result<CompletionResponse> chatCompletion(@RequestBody ChatCompletionDTO chatCompletion) {
        Result<CompletionResponse> result=new Result<>();
        CompletionResponse completionResponse=openAIService.chatCompletion(chatCompletion);
        if (completionResponse!=null) {
            return result.ok(completionResponse);
        }
        return result.fail(null,-1,"失败");
    }

    @RequestMapping(value = "/chat/unify/completion", method = RequestMethod.POST, produces = {MediaType.APPLICATION_JSON_VALUE, MediaType.TEXT_EVENT_STREAM_VALUE})
    public Object unifyChatCompletion(@RequestBody UnifyChatCompletion chatCompletion, HttpServletResponse response) {
        Result<UnifyCompletion> result = new Result<>();
        if (!chatCompletion.getStream()) {
            try {
                UnifyCompletion completionResponse = openAIService.unifyChatCompletion(chatCompletion);
                if (completionResponse != null) {
                    return ResponseEntity.ok(result.ok(completionResponse));
                }
            } catch (Exception ex) {
                return new ResponseEntity<>(new Result<>().fail(null, -1, ex.getMessage()), HttpStatus.OK);
            }
        } else {
            try {
                SseEmitter emitter = new SseEmitter(60000L * 5);
                executorService.execute(() -> {
                    try {
                        InputStream inputStream = openAIService.unifyChatCompletionStream(chatCompletion);
                        InputStreamReader reader = new InputStreamReader(inputStream);
                        BufferedReader bufferedReader = new BufferedReader(reader);
                        String line;
                        try {
                            while ((line = bufferedReader.readLine()) != null) {
                                log.info("流式消息:"+line);
                                emitter.send(SseEmitter.event().data(line));
                                UnifyStreamMessage unifyStreamMessage=  JSON.parseObject(line, UnifyStreamMessage.class);
                                if (unifyStreamMessage!=null&& !unifyStreamMessage.getChoices().isEmpty()){
                                    boolean isEnd=false;
                                    for (UnifyChoice unifyChoice:unifyStreamMessage.getChoices()){
                                        if ("stop".equals(unifyChoice.getFinishReason())){
                                            isEnd=true;
                                            break;
                                        }
                                    }
                                    if (isEnd) {
                                        break;
                                    }
                                }
                            }
                        }catch (Exception exception){
                            exception.printStackTrace();
                        }
                        inputStream.close();
                        emitter.complete();
                        log.info("流式消息结束");
                    } catch (IOException e) {
                        emitter.completeWithError(e);
                    }
                });
                return emitter;
            } catch (Exception ex) {
                return new ResponseEntity<>(new Result<>().fail(null, -1, ex.getMessage()), HttpStatus.OK);
            }
        }
        return new ResponseEntity<>(HttpStatus.NOT_FOUND);
    }

    private String getBody(HttpServletRequest request) {
        try (InputStream inputStream = request.getInputStream()) {
            return IOUtils.toString(inputStream, StandardCharsets.UTF_8);
        } catch (IOException e) {
            log.error("get input error", e);
        }
        return null;
    }


}
